from model import *
from types import SimpleNamespace
from dataset import *
from torch.utils.data import DataLoader
from utils import *
import pickle

device='cuda'
batch_size=128
maxlen=50
sample=0
epoch=220

args={}
data="Sports"
args['device']=device
args['hidden_units']=50
args['maxlen']=maxlen
args['dropout_rate']=0.2
args['num_blocks']=2
args['num_heads']=1
model_path=f"/mnt/data/users/fyb/mywork/DLRec/result/Sports/sasrec/SASRec.epoch=240.lr=0.001.layer=2.head=1.hidden=50.maxlen=50.pth"
args = SimpleNamespace(**args)
dataset = SeqDataset("data/"+data,maxlen=maxlen)
candidate_dataloader=DataLoader(dataset.candidate_item,batch_size=batch_size,collate_fn=collate_fn3)
itemnum=dataset.item_max-1
sasmodel = SASRec(1,itemnum, args).to(args.device)
for name, param in sasmodel.named_parameters():
    try:
        torch.nn.init.xavier_normal_(param.data)
    except:
        pass  # just ignore those failed init layers

sasmodel.pos_emb.weight.data[0, :] = 0
sasmodel.item_emb.weight.data[0, :] = 0
sasmodel.load_state_dict(torch.load(model_path, map_location=torch.device(args.device)))

pickle.dump(sasmodel.item_emb.weight.detach(),open(f'./embedding/{data}_embedding.pkl', 'wb'))
